import matplotlib.pyplot as plt
import numpy as np
from sklearn.decomposition import PCA

from view.endSrc.UnionFindSet import UnionFindSet


class MergeCentersByMultiOb:
    '''
        add by ChenHongjie
        return mergedCentersData
    '''

    def __init__(self, dataset, oblst, isPlot=False):
        self.dataset = dataset
        self.oblst = oblst
        self.componentLabels_ob = [ob.getComponentLabels() for ob in self.oblst]
        self.obNum = len(self.oblst)
        self.isPlot = isPlot

        if self.isPlot:
            if dataset.shape[1] > 2:
                pca = PCA(n_components=2)
                self.dataset2d = pca.fit_transform(dataset)
            else:
                self.dataset2d = dataset

    def merge(self, cidlst, centersFromObs, componentCenterNum_ob):
        if self.obNum < 2:
            mergedCentersData = self.dataset[cidlst, :].tolist()
            mergedCentersData = np.unique(mergedCentersData, axis=0).tolist()
            return mergedCentersData

        # plot
        if self.isPlot:
            plt.scatter(self.dataset2d[:, 0], self.dataset2d[:, 1])
            plt.scatter(self.dataset2d[cidlst, 0], self.dataset2d[cidlst, 1], c='r')

        ufs = UnionFindSet(cidlst)
        centersNum = len(cidlst)

        for i in range(centersNum):
            for j in range(i+1, centersNum):
                mergeScore = 0
                id1 = cidlst[i]  # center point id
                id2 = cidlst[j]
                obIdx1 = centersFromObs[i]  # from which observer
                obIdx2 = centersFromObs[j]

                # if two centers are found by "the same observer" in the "same" component, tend NOT to merge
                if (obIdx1 == obIdx2) and (self.componentLabels_ob[obIdx1][id1] == self.componentLabels_ob[obIdx2][id2]):
                    mergeScore -= 1

                for obidx, componentLabels in enumerate(self.componentLabels_ob):
                    compIdx1 = componentLabels[id1]
                    compIdx2 = componentLabels[id2]
                    if compIdx1 == compIdx2:
                        if componentCenterNum_ob[obidx][compIdx2] == 1:
                            mergeScore += 0.5
                        mergeScore += 1
                    elif (abs(compIdx2 - compIdx1) == 1) and \
                            (componentCenterNum_ob[obidx][compIdx1] > 1 and componentCenterNum_ob[obidx][compIdx2] > 1):
                        mergeScore += (self.obNum / 2.0 - 1) / (self.obNum - 1)

                if mergeScore >= self.obNum / 2.0:
                    ufs.union(id1, id2)
                    # plot edge
                    if self.isPlot:
                        plt.plot([self.dataset2d[id1, 0], self.dataset2d[id2, 0]], [self.dataset2d[id1, 1], self.dataset2d[id2, 1]], 'k')
        if self.isPlot:
            plt.show()

        merged_cidlst = ufs.getAllHead()
        merged_cidlst.sort()

        mergedCentersData = self.dataset[merged_cidlst, :].tolist()
        mergedCentersData = np.unique(mergedCentersData, axis=0).tolist()
        return mergedCentersData
